import PIL
import numpy as np
import torch
import random
import itertools
import rootutils
import sys
sys.path.append('/lc/code/3D/a3R/src')
rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
from src.dust3r.datasets.base.easy_dataset import EasyDataset
from src.dust3r.datasets.utils.transforms import ImgNorm, SeqColorJitter
from src.dust3r.utils.geometry import depthmap_to_absolute_camera_coordinates
import src.dust3r.datasets.utils.cropping as cropping
from src.dust3r.datasets.utils.corr import extract_correspondences_from_pts3d
from src.dust3r.datasets.hypersim import HyperSim_Multi
from src.dust3r.datasets.wildrgbd2 import WildRGBD_Multi
import argparse

def get_parser():
    parser = argparse.ArgumentParser(
        description="Preprocess the BlendedMVS dataset by converting camera parameters, "
        "resizing images, and saving processed outputs in Hypersim format."
    )
    parser.add_argument(
        "--wildrgbd_dir", 
        default='processed',
        help="Root directory of the BlendedMVS dataset."
    )
    return parser
if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    ds = WildRGBD_Multi(allow_repeat=False, split='train', ROOT=f'/lc/data/3D/wildrgbd/{args.wildrgbd_dir}', aug_crop=0, resolution=[(224, 224)], num_views=3, n_corres=0, samples_per_scene=1, check=True)
    print(ds)
    for mask_bg in [True, False, "rand"]:
        ds.check_depth_valid(actual_mask_bg=mask_bg)

    